"""计划执行"""
import json
import asyncio
import time
import uuid
from typing import Any, Dict, List, Optional
from langgraph.store.memory import BaseStore
from json_repair import repair_json

from .planner import Plan, PlanStep
from .tool_manager import ToolManager
from model.excutor_model import StepResult, ExecutionStatus, ExecutionResult
from model.context_model import ToolContext
from model.memory_model import MemoryType
from loguru import logger
import traceback

class PlanExecutor:
    """计划执行"""
    
    def __init__(self, memory_store: BaseStore, tool_manager: ToolManager):
        self.tool_manager = tool_manager
        self.memory_store = memory_store
        self.execution_history: List[ExecutionResult] = []

    async def execute_plan(self, plan: Plan, max_retries: int = 2) -> ExecutionResult:
        """执行一个完整计划"""
        start_time = time.time()
        step_results = []
        completed_steps = []
        failed_steps = []

        execution_context = {
            "plan_id": plan.id,
            "query": plan.query,
            "variables": {},
            "step_outputs": {}
        }
        
        try:
            if plan.plan_type.value == "sequential":
                step_results = await self._execute_sequential(plan, execution_context, max_retries)
            elif plan.plan_type.value == "parallel":
                step_results = await self._execute_parallel(plan, execution_context, max_retries)
            elif plan.plan_type.value == "conditional":
                step_results = await self._execute_conditional(plan, execution_context, max_retries)
            elif plan.plan_type.value == "iterative":
                step_results = await self._execute_iterative(plan, execution_context, max_retries)
            else:
                step_results = await self._execute_sequential(plan, execution_context, max_retries)
            
            # 分析执行结果
            completed_steps = [r.step_id for r in step_results if r.status == ExecutionStatus.COMPLETED]
            failed_steps = [r.step_id for r in step_results if r.status == ExecutionStatus.FAILED]
            
            # 执行状态更新
            if len(completed_steps) == len(plan.steps):
                overall_status = ExecutionStatus.COMPLETED
            elif len(completed_steps) > 0:
                overall_status = ExecutionStatus.RUNNING
            else:
                overall_status = ExecutionStatus.FAILED
            
            # 计划完成率
            success_rate = len(completed_steps) / len(plan.steps) if plan.steps else 0.0
            
            # 计划执行结果
            final_output = self._determine_final_output(step_results, execution_context)
            
        except Exception as e:
            logger.error(traceback.format_exc())
            overall_status = ExecutionStatus.FAILED
            success_rate = 0.0
            final_output = f"Execution failed: {str(e)}"

        execution_result = ExecutionResult(
            plan_id=plan.id,
            status=overall_status,
            step_results=step_results,
            total_time=time.time() - start_time,
            success_rate=success_rate,
            final_output=final_output,
            metadata={
                "completed_steps": completed_steps,
                "failed_steps": failed_steps,
                "plan_type": plan.plan_type.value
            }
        )
        
        # 保存计划执行情况
        self.execution_history.append(execution_result)
        
        return execution_result
    
    async def _execute_sequential(self, plan: Plan, context: Dict[str, Any], 
                                 max_retries: int) -> List[StepResult]:
        """顺序执行计划"""
        results = []
        
        for step in plan.steps:
            if not self._dependencies_satisfied(step, results):
                results.append(StepResult(
                    step_id=step.id,
                    status=ExecutionStatus.SKIPPED,
                    output=None,
                    error="不满足依赖条件"
                ))
                continue

            result = await self._execute_step(step, context, max_retries)
            results.append(result)
            
            # 步骤完成后，更新上下文
            if result.status == ExecutionStatus.COMPLETED:
                context["step_outputs"][step.id] = result.output
                context["variables"][f"step_{step.id}_output"] = result.output
            
            # 若关键步骤失败，则停止。
            if result.status == ExecutionStatus.FAILED and step.confidence > 0.8:
                break
        
        return results
    
    async def _execute_parallel(self, plan: Plan, context: Dict[str, Any], 
                               max_retries: int) -> List[StepResult]:
        """并行执行计划"""
        results = []
        remaining_steps = plan.steps.copy()
        
        while remaining_steps:
            # 查找满足依赖条件的可执行步骤
            executable_steps = []
            for step in remaining_steps:
                if self._dependencies_satisfied(step, results):
                    executable_steps.append(step)
            
            if not executable_steps:
                # 若没有更多可执行的步骤，则将剩余步骤标记为已跳过
                for step in remaining_steps:
                    results.append(StepResult(
                        step_id=step.id,
                        status=ExecutionStatus.SKIPPED,
                        output=None,
                        error="不满足依赖条件"
                    ))
                break
            
            # 并行执行
            tasks = []
            for step in executable_steps:
                task = self._execute_step(step, context.copy(), max_retries)
                tasks.append(task)
            step_results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 结果处理
            for i, result in enumerate(step_results):
                if isinstance(result, Exception):
                    result = StepResult(
                        step_id=executable_steps[i].id,
                        status=ExecutionStatus.FAILED,
                        output=None,
                        error=str(result)
                    )
                
                results.append(result)
                
                # 更新上下文
                if result.status == ExecutionStatus.COMPLETED:
                    context["step_outputs"][result.step_id] = result.output
                    context["variables"][f"step_{result.step_id}_output"] = result.output
            
            # 移除已执行步骤
            executed_ids = {step.id for step in executable_steps}
            remaining_steps = [s for s in remaining_steps if s.id not in executed_ids]
        
        return results
    
    async def _execute_conditional(self, plan: Plan, context: Dict[str, Any], 
                                  max_retries: int) -> List[StepResult]:
        """条件执行计划"""
        results = []
        
        for step in plan.steps:
            # 检查是否满足执行条件
            if not self._should_execute_step(step, context, results):
                results.append(StepResult(
                    step_id=step.id,
                    status=ExecutionStatus.SKIPPED,
                    output=None,
                    error="条件不满足"
                ))
                continue
            
            # 依赖检查
            if not self._dependencies_satisfied(step, results):
                results.append(StepResult(
                    step_id=step.id,
                    status=ExecutionStatus.SKIPPED,
                    output=None,
                    error="不满足依赖关系"
                ))
                continue
            
            # 执行
            result = await self._execute_step(step, context, max_retries)
            results.append(result)
            
            # 更新上下文
            if result.status == ExecutionStatus.COMPLETED:
                context["step_outputs"][step.id] = result.output
                context["variables"][f"step_{step.id}_output"] = result.output
        
        return results
    
    async def _execute_iterative(self, plan: Plan, context: Dict[str, Any], 
                                max_retries: int, max_iterations: int = 5) -> List[StepResult]:
        """迭代执行步骤，直到满足条件为止。"""
        all_results = []
        iteration = 0
        
        while iteration < max_iterations:
            iteration_results = []
            
            for step in plan.steps:
                # 检查依赖
                if not self._dependencies_satisfied(step, iteration_results):
                    continue
                
                # 执行步骤
                result = await self._execute_step(step, context, max_retries)
                iteration_results.append(result)
                
                # 更新上下文
                if result.status == ExecutionStatus.COMPLETED:
                    context["step_outputs"][f"{step.id}_iter_{iteration}"] = result.output
                    context["variables"][f"step_{step.id}_output"] = result.output
            
            all_results.extend(iteration_results)
            
            # 判断是否需要继续迭代
            if not self._should_continue_iteration(iteration_results, context):
                break
            iteration += 1
        
        return all_results
    
    async def _execute_step(self, step: PlanStep, context: Dict[str, Any], 
                           max_retries: int) -> StepResult:
        """单步执行."""
        start_time = time.time()
        
        for attempt in range(max_retries + 1):
            try:
                # 通过替换变量来准备输入
                step_input = self._substitute_variables(step.input_template, context)
                
                # 工具调用
                tool_result = await self.tool_manager.execute_tool(step.tool, step_input)
                
                # 创建工具调用的上下文信息
                tool_context = ToolContext(
                    tool_name=step.tool,
                    input_data=step_input,
                    output_data=tool_result.data if tool_result.success else None,
                    success=tool_result.success,
                    error_message=tool_result.error,
                    execution_time=time.time() - start_time
                )
                
                # 添加到上下文管理器
                self.add_tool_context(tool_context)

                if tool_result.success:
                    return StepResult(
                        step_id=step.id,
                        status=ExecutionStatus.COMPLETED,
                        output=tool_result.data,
                        execution_time=time.time() - start_time,
                        metadata={"attempts": attempt + 1}
                    )
                else:
                    if attempt == max_retries:
                        return StepResult(
                            step_id=step.id,
                            status=ExecutionStatus.FAILED,
                            output=None,
                            error=tool_result.error,
                            execution_time=time.time() - start_time,
                            metadata={"attempts": attempt + 1}
                        )
                    await asyncio.sleep(0.1)
            
            except Exception as e:
                if attempt == max_retries:
                    return StepResult(
                        step_id=step.id,
                        status=ExecutionStatus.FAILED,
                        output=None,
                        error=str(e),
                        execution_time=time.time() - start_time,
                        metadata={"attempts": attempt + 1}
                    )
                await asyncio.sleep(1)

        return StepResult(
            step_id=step.id,
            status=ExecutionStatus.FAILED,
            output=None,
            error="Unexpected execution path",
            execution_time=time.time() - start_time
        )
    
    def _dependencies_satisfied(self, step: PlanStep, completed_results: List[StepResult]) -> bool:
        """检查是否满足依赖条件"""
        completed_step_ids = {r.step_id for r in completed_results if r.status == ExecutionStatus.COMPLETED}
        return all(dep in completed_step_ids for dep in step.dependencies)
    
    def _should_execute_step(self, step: PlanStep, context: Dict[str, Any], results: List[StepResult]) -> bool:
        """检查执行所需条件"""
        if not step.conditions:
            return True

        # 这里仅检查输入变量
        for condition_key, condition_value in step.conditions.items():
            if condition_key in context["variables"]:
                if context["variables"][condition_key] != condition_value:
                    return False
        return True
    
    def _should_continue_iteration(self, iteration_results: List[StepResult], context: Dict[str, Any]) -> bool:
        """判断是否需要继续迭代"""
        # 简单的启发式规则：若有任何步骤失败，则继续迭代执行。
        return any(r.status == ExecutionStatus.FAILED for r in iteration_results)
    
    def _substitute_variables(self, template: str, context: Dict[str, Any]) -> Dict:
        """替换输入模板中的变量。"""
        result = template
        for key, value in context.get("variables", {}).items():
            placeholder = f"{{{key}}}"
            if placeholder in result:
                result = result.replace(placeholder, json.dumps(value, ensure_ascii=False))

        for step_id, output in context.get("step_outputs", {}).items():
            placeholder = f"{{step_{step_id}_output}}"
            if placeholder in result:
                result = result.replace(placeholder, json.dumps(output, ensure_ascii=False))
        return repair_json(result, return_objects=True)
    
    def _determine_final_output(self, step_results: List[StepResult], context: Dict[str, Any]) -> Any:
        """根据步骤结果确定最终输出。"""
        for result in reversed(step_results):
            if result.status == ExecutionStatus.COMPLETED and result.output is not None:
                return result.output
        return {
            "message": "计划执行完成，部分步骤执行失败. ",
            "successful_steps": len([r for r in step_results if r.status == ExecutionStatus.COMPLETED]),
            "total_steps": len(step_results)
        }
    
    def get_execution_stats(self) -> Dict[str, Any]:
        if not self.execution_history:
            return {"total_executions": 0}
        
        total_executions = len(self.execution_history)
        successful_executions = len([e for e in self.execution_history if e.status == ExecutionStatus.COMPLETED])
        avg_success_rate = sum(e.success_rate for e in self.execution_history) / total_executions
        avg_execution_time = sum(e.total_time for e in self.execution_history) / total_executions
        
        return {
            "total_executions": total_executions,
            "successful_executions": successful_executions,
            "success_rate": successful_executions / total_executions,
            "avg_step_success_rate": avg_success_rate,
            "avg_execution_time": avg_execution_time
        }

    def add_tool_context(self, tool_context: ToolContext):
        """添加工具执行的上下文."""
        self.memory_store.put(
            (MemoryType.TOOL_CONTEXT.value, ),
            key=str(uuid.uuid4()),
            value={
                'name': tool_context.tool_name,
                'tool_context': tool_context.model_dump_json()
            },
            index=['name']
        )